import os
from math import ceil
from argparse import ArgumentParser

from tqdm import tqdm
from affine import Affine
import rasterio
from rasterio.windows import Window



parser = ArgumentParser()
parser.add_argument(
    '-p', '--path', help='the folder under which the files need mosaiced'
)
parser.add_argument(
    '-s', '--save_dir', help='the folder to save the splited files'
)
parser.add_argument(
    '-f', '--prefix', default='', help='prefix of out fileName'
)
parser.add_argument(
    '-k', '--kwt', default='', help='key word (means date) of spacenet in output file name'
)
args = parser.parse_args()

#--------- Param list
split_size = 900
split_stride = split_size - 320

kwt = args.kwt
prefix = args.prefix
path = args.path
save_dir = args.save_dir



src_raster = rasterio.open(path)
geotrans, height, width = list(src_raster.transform), src_raster.height, src_raster.width
res_x, res_y = geotrans[0], -geotrans[4]
min_word_x, max_word_y = geotrans[2], geotrans[5]
kwargs = src_raster.meta.copy()

start_x = min_word_x - (split_size-split_stride) * res_x
end_x = start_x + width * res_x
start_y = max_word_y + (split_size-split_stride) * res_y
end_y = start_y - height * res_y
stride_x = split_stride * res_x
stride_y = split_stride * res_y

cur_y = start_y
cnt = 1

# pbar = tqdm(
#     total=(int(start_y-end_y)//stride_y+1)*\
#         (int(end_x-start_x)//stride_x+1), 
#     ncols=100
# )
with tqdm(total=(int(start_y-end_y)//stride_y+1)*\
                (int(end_x-start_x)//stride_x+1)) as pbar:
    while cur_y >= end_y:
        cur_x = start_x

        while cur_x <= end_x:
            
            start_col = ceil((cur_x-min_word_x) / res_x)
            end_col = start_col + split_size
            start_row = ceil((max_word_y-cur_y) / res_y)
            end_row = start_row + split_size

            src_win_rows_start = start_row if start_row>=0 else 0
            dst_win_rows_start = 0 if start_row>=0 else -start_row
            src_win_cols_start = start_col if start_col>=0 else 0
            dst_win_cols_start = 0 if start_col>=0 else -start_col

            src_win_rows_end = end_row if end_row<=height else height
            dst_win_rows_end = split_size if end_row<=height else split_size-(end_row-height)
            src_win_cols_end = end_col if end_col<=width else width
            dst_win_cols_end = split_size if end_col<=width else split_size-(end_col-width)

            src_win = Window.from_slices(
                (src_win_rows_start, src_win_rows_end), 
                (src_win_cols_start, src_win_cols_end)
            )
            dst_win = Window.from_slices(
                (dst_win_rows_start, dst_win_rows_end), 
                (dst_win_cols_start, dst_win_cols_end)
            )

            filename = "{}_{}_{}_tile_{:0>4}.tif".format(prefix, kwt, kwt, cnt)
            kwargs.update({
                'height': split_size,
                'width': split_size,
                'transform': Affine.from_gdal(*(cur_x, res_x, 0, cur_y, 0, -res_y))
            })
            dst_raster = rasterio.open(os.path.join(save_dir, filename), mode='w+', **kwargs)

            src_data = src_raster.read(window=src_win, masked=True)
            dst_raster.write(src_data, window=dst_win)
            dst_mask = dst_raster.read(masked=True).mask
            if len(dst_mask.shape) > 2:
                dst_mask = dst_mask[0]

            dst_raster.close()
            thresh = (2*split_stride-split_size) * (2*split_stride-split_size)
            if (dst_mask[
                    split_size-split_stride: split_stride, 
                    split_size-split_stride: split_stride
                    ].sum() == thresh
                ):
                os.remove(os.path.join(save_dir, filename))
                cnt -= 1

            cur_x += stride_x
            cnt += 1
            pbar.update(1)

        cur_y -= stride_y
